#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/12/31 下午5:41
# @Author  : HUANG Xiong
# @Site    : 
# @File    : macro_value_related.py
# @Software: PyCharm 

"""
脚本说明：获取宏观指标数据
"""
from quant_researcher.quant.project_tool.db_operator import db_conn, my_mysql
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.celebrity import get_lst_condition
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.datasource_fetch.macro_api import macro_info

T_INDIC_VALUE = 'indicator_val'
T_INDIC_VALUE_JY = 'indicator_data'
EARLIEST_START_DATE = '1990-01-01'
EARLIEST_START_DATE_JY = '19900101'


def get_indicator_value(indic_code=None, indic_name=None,
                        start_date=None, end_date=None, only_latest=False):
    """
    获得宏观指标的数据，数据源为聚源

    :param str or list indic_code: 指标代码，支持单个（str）、多个（list）
    :param str or list indic_name: 指标名称，可以根据指标名称查询，但是注意指标名称必须非常准确，否则无法查出
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :return: pd.DataFrame
                +---------------+------------+-------------+
                |   indic_code  |  end_date  | indic_value |
                +---------------+------------+-------------+
                | 2011040000002 | 1990-12-19 |    99.98    |
                | 2011040000002 | 1990-12-20 |    104.39   |
                | 2011040000002 | 1990-12-21 |    109.13   |
                | 2011040000002 | 1990-12-24 |    114.55   |
                | 2011040000002 | 1990-12-25 |    120.25   |
                +---------------+------------+-------------+
    """
    # 只有指标代码没有输入，但是指标名称不为None的情况下，才根据指标名称进行查询
    if indic_code is None and indic_name is not None:
        indic_code = macro_info.get_indicator_code_by_name(indic_name=indic_name)
        if indic_code is None:
            LOG.error(f"没有找到{indic_name}的指标代码信息，请检查该指标是否存在")

    # select = ['indic_code', 'end_date', 'data_value * power(10, power_number) as indic_value']
    select = ['indic_code', 'end_date', 'data_value as indic_value']

    where = []

    # indic_code
    if isinstance(indic_code, str) or isinstance(indic_code, list):
        if isinstance(indic_code, str):
            indic_code = [indic_code]
        indic_where_condition = get_lst_condition(indic_code, 'indic_code')
        where.append(indic_where_condition)
    else:
        LOG.error('目前indic_code只支持str、list，请按要求输入')
        raise NotImplementedError

    # 将日期转换成带杠字符串
    if start_date is None:
        start_date = EARLIEST_START_DATE_JY
    else:
        start_date = time_tool.format_date_str(start_date, '%Y%m%d')
    if end_date is None:
        end_date = time_tool.get_today(marker='without_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date, '%Y%m%d')

    conn = db_conn.get_stock_conn()
    # 只返回最新数据
    if only_latest:
        select = ', '.join(select)
        sql_str = (f"select {select} from {T_INDIC_VALUE_JY} as a, "
                   f"(select indic_code as code, max(end_date) as enddate "
                   f"from {T_INDIC_VALUE_JY} "
                   f"where end_date <= '{end_date}' "
                   f"and {indic_where_condition} "
                   f"group by indic_code) as b "
                   f"where a.indic_code = b.code "
                   f"and a.end_date = b.enddate ")
        df = my_mysql.read_v3(sql_str, conn)
    # 返回一段时间数据的where语句
    else:
        where.append(f"end_date >= '{start_date}'")
        where.append(f"end_date <= '{end_date}'")
        df = my_mysql.read_v2(select=select, where=where, sfrom=T_INDIC_VALUE_JY, conn=conn)

    conn.close()

    if df.empty:
        LOG.error(f"没有找到{indic_code}指数的数据，请检查")
        return

    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


def get_indicator_value_and_info(indic_code=None, indic_name=None, start_date=None,
                                 end_date=None, only_latest=False):
    """
    将指标的值和信息合并

    :param str or list indic_code: 指标代码，支持单个（str）、多个（list）
    :param str or list indic_name: 指标名称，可以根据指标名称查询，但是注意指标名称必须非常准确，否则无法查出
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :return:
    """
    if indic_code is None and indic_name is None:
        LOG.error(f"即没有输入{indic_code}也没有输入{indic_name}，请至少输入一项进行搜索")
        return

    if indic_code is None and indic_name is not None:  # 需要根据指标名称获取指标代码
        indic_code = macro_info.get_indicator_code_by_name(indic_name=indic_name)
        if indic_code is None:
            LOG.error(f"没有找到{indic_name}的指标代码信息，请检查该指标是否存在")
            return

    value_df = get_indicator_value(indic_code, indic_name, start_date, end_date, only_latest)
    if value_df is None:
        return
    value_df = value_df.dropna()
    if value_df.empty:
        return

    info_df = macro_info.get_indicator_info(
        indic_code, indic_name, select=['indic_code', 'indic_name', 'frequency', 'unit', 'unit_report', 'indic_state'])

    if info_df is None:
        return
    df = value_df.merge(info_df, how='inner', on='indic_code')
    if df.empty:
        return
    df['indic_code'] = df['indic_code'].astype(str)
    return df


def get_indicator_value_old(indic_code=None, indic_name=None, start_date=None, end_date=None,
                        only_latest=False, select=None):
    """
    获得宏观指标的数据（数据来源为爬取萝卜投研网站，已弃用，目前数据源为聚源）

    :param str or list indic_code: 指标代码，支持单个（str）、多个（list）
    :param str or list indic_name: 指标名称，可以根据指标名称查询，但是注意指标名称必须非常准确，否则无法查出
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :param list select: 需要查询哪些字段, 不指定的话默认['indic_code', 'val_date as end_date', 'data_value as indic_value']
    :return: pd.DataFrame
                +---------------+------------+-------------+
                |   indic_code  |  end_date  | indic_value |
                +---------------+------------+-------------+
                | 2011040000002 | 1990-12-19 |    99.98    |
                | 2011040000002 | 1990-12-20 |    104.39   |
                | 2011040000002 | 1990-12-21 |    109.13   |
                | 2011040000002 | 1990-12-24 |    114.55   |
                | 2011040000002 | 1990-12-25 |    120.25   |
                +---------------+------------+-------------+
    """
    # 只有指标代码没有输入，但是指标名称不为None的情况下，才根据指标名称进行查询
    if indic_code is None and indic_name is not None:
        indic_code = macro_info.get_indicator_code_by_name(indic_name=indic_name)
        if indic_code is None:
            LOG.error(f"没有找到{indic_name}的指标代码信息，请检查该指标是否存在")

    if select is None:
        select = ['indic_code', 'val_date as end_date', 'data_value as indic_value']

    where = []

    # indic_code
    if isinstance(indic_code, str) or isinstance(indic_code, list):
        if isinstance(indic_code, str):
            indic_code = [indic_code]
        indic_where_condition = get_lst_condition(indic_code, 'indic_code')
        where.append(indic_where_condition)
    else:
        LOG.error('目前indic_code只支持str、list，请按要求输入')
        raise NotImplementedError

    # 将日期转换成带杠字符串
    if start_date is None:
        start_date = EARLIEST_START_DATE
    else:
        start_date = time_tool.format_date_str(start_date)
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date)

    conn = db_conn.get_stock_conn()
    # 只返回最新数据
    if only_latest:
        select = ', '.join(select)
        sql_str = (f"select {select} from {T_INDIC_VALUE} as a, "
                   f"(select indic_code as code, max(val_date) as enddate "
                   f"from {T_INDIC_VALUE} "
                   f"where val_date <= '{end_date}' "
                   f"and {indic_where_condition} "
                   f"group by indic_code) as b "
                   f"where a.indic_code = b.code "
                   f"and a.val_date = b.enddate ")
        df = my_mysql.read_v3(sql_str, conn)
    # 返回一段时间数据的where语句
    else:
        where.append(f"val_date >= '{start_date}'")
        where.append(f"val_date <= '{end_date}'")
        df = my_mysql.read_v2(select=select, where=where, sfrom=T_INDIC_VALUE, conn=conn)

    conn.close()

    if df.empty:
        LOG.error(f"没有找到{indic_code}指数的数据，请检查")
        return

    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


def get_indicator_value_and_info_old(indic_code=None, indic_name=None, start_date=None,
                                 end_date=None, only_latest=False):
    """
    将指标的值和信息合并（数据来源为爬取萝卜投研网站，已弃用，目前数据源为聚源）

    :param str or list indic_code: 指标代码，支持单个（str）、多个（list）
    :param str or list indic_name: 指标名称，可以根据指标名称查询，但是注意指标名称必须非常准确，否则无法查出
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :return:
    """
    if indic_code is None and indic_name is None:
        LOG.error(f"即没有输入{indic_code}也没有输入{indic_name}，请至少输入一项进行搜索")
        return

    if indic_code is None and indic_name is not None:  # 需要根据指标名称获取指标代码
        indic_code = macro_info.get_indicator_code_by_name_old(indic_name=indic_name)
        if indic_code is None:
            LOG.error(f"没有找到{indic_name}的指标代码信息，请检查该指标是否存在")
            return

    value_df = get_indicator_value_old(indic_code, indic_name, start_date, end_date, only_latest)
    if value_df is None:
        return
    value_df = value_df.dropna()
    if value_df.empty:
        return

    info_df = macro_info.get_indicator_info_old(indic_code, indic_name, select=['indic_code', 'indic_name',
                                                                            'frequency', 'unit'])

    if info_df is None:
        return
    df = value_df.merge(info_df, how='inner', on='indic_code')
    if df.empty:
        return
    # 指标代码变成字符串格式
    df['indic_code'] = df['indic_code'].astype(str)

    return df


if __name__ == '__main__':
    ccc = get_indicator_value_and_info(indic_name='居民居住消费价格指数CPI_(上年=100)_当月')
    aaa = get_indicator_value('1000110002779', only_latest=False)
